package com.asgard.keybox;

import java.io.File;
import java.io.IOException;
import java.io.InputStream;
import java.math.BigInteger;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.cert.CertificateException;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.stream.Collectors;

import javax.xml.bind.JAXBContext;
import javax.xml.bind.Marshaller;
import javax.xml.bind.Unmarshaller;

import org.apache.commons.lang3.Validate;
import org.apache.commons.lang3.tuple.Pair;
import org.bouncycastle.util.encoders.Hex;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.asgard.cert.Certificate;
import com.asgard.crypto.CryptoService;
import com.asgard.crypto.ECKey;
import com.asgard.keybox.KeyBoxData.CertificateData;
import com.asgard.keybox.KeyBoxData.ECKeyData;
import com.asgard.utils.CryptoUtils;
import com.google.common.collect.Ordering;

/**
 * 该类用于本地密钥与证书信息管理。
 */
public class KeyBox {

    private static final Logger log = LoggerFactory.getLogger(KeyBox.class);

    /**
     * 密钥存储项
     */
    private static class ECKeyItem {
        String name;
        ECKeyData keyData;
        ECKey key;
    }

    /**
     * 证书存储项
     */
    private static class CertItem {
        String name;
        CertificateData certData;
        Certificate cert;
        PrivateKey privKey;
    }

    /**
     * 密钥列表: 名称 -> 密钥实例
     */
    private Map<String, ECKeyItem> keys;
    /**
     * 证书列表: 名称 -> 证书实例
     */
    private Map<String, CertItem> certs;
    /**
     * 证书列表: <MEMBER_ID,SN> -> 证书实例
     */
    private Map<Pair<String, BigInteger>, CertItem> memberIdSnCerts;

    /**
     * 加密密钥
     */
    private byte[] aesKey;
    /**
     * 盐
     */
    private byte[] salt;

    /**
     * 初始化密钥库。
     */
    public KeyBox() {
        keys = new ConcurrentHashMap<>();
        certs = new ConcurrentHashMap<>();
        memberIdSnCerts = new ConcurrentHashMap<>();
    }

    /**
     * 初始化新的密钥库。
     *
     * @param pwd 初始密码
     */
    public void initNew(String pwd) {
        salt = CryptoService.get().generateKey();
        aesKey = CryptoService.get().generateKey(pwd, salt, false);
    }

    /**
     * 从文件加载密钥库数据。
     *
     * @param file     文件
     * @param password 密码
     * @throws IOException 输入输出异常
     */
    public void loadFromInputstream(InputStream is, String password) throws IOException {
        try {
            // 读取文件
            log.info("读取密钥库文件");
            // 将XML反序列化为对象
            JAXBContext context = JAXBContext.newInstance(KeyBoxData.class);
            Unmarshaller unmarshaller = context.createUnmarshaller();
            Object obj = unmarshaller.unmarshal(is);

            loadKeyBox(password, obj);
        } catch (Exception e) {
            throw new IOException(e);
        }
    }

    /**
     * 从文件加载密钥库数据。
     *
     * @param file     文件
     * @param password 密码
     * @throws IOException 输入输出异常
     */
    public void loadFromFile(File file, String password) throws IOException {
        try {
            // 读取文件
            log.info("读取密钥库文件: {}", file.getAbsolutePath());
            // 将XML反序列化为对象
            JAXBContext context = JAXBContext.newInstance(KeyBoxData.class);
            Unmarshaller unmarshaller = context.createUnmarshaller();
            Object obj = unmarshaller.unmarshal(file);

            loadKeyBox(password, obj);
        } catch (Exception e) {
            throw new IOException(e);
        }
    }

    private void loadKeyBox(String password, Object obj) throws IOException, CertificateException {
        if (obj == null || !(obj instanceof KeyBoxData))
            throw new IOException("读取XML文件错误");
        KeyBoxData keyBoxData = (KeyBoxData) obj;
        // 读取Salt
        Validate.notNull(keyBoxData.getSalt(), "盐值不得为空");
        this.salt = Hex.decode(keyBoxData.getSalt());
        // 初始化加密工具
        this.aesKey = CryptoService.get().generateKey(password, salt, false);
        // 读取密钥
        for (ECKeyData keyData : keyBoxData.getKeys()) {
            Validate.notNull(keyData.getName(), "密钥名称不得为空");
            Validate.notNull(keyData.getPublicKey(), "公钥不得为空");
            ECKeyItem key = new ECKeyItem();
            key.name = keyData.getName();
            key.keyData = keyData;
            PublicKey pubKey = CryptoService.get().parsePublicKey(Hex.decode(keyData.getPublicKey()));
            PrivateKey privKey = null;
            if (keyData.getPrivateKey() != null && !keyData.getPrivateKey().isEmpty()) {
                byte[] privKeyBytes = CryptoService.get().symmetricDecrypt(aesKey, Hex.decode(keyData.getPrivateKey()));
                privKey = CryptoService.get().parsePrivateKey(privKeyBytes);
                Arrays.fill(privKeyBytes, (byte) 0);
            }
            key.key = new ECKey(privKey, pubKey);
            // 将密钥添加到列表中
            keys.put(key.name, key);
        }

        // 读取证书
        for (CertificateData certData : keyBoxData.getCertificates()) {
            Validate.notNull(certData.getName(), "证书名称不得为空");
            Validate.notNull(certData.getContent(), "证书内容不得为空");
            CertItem cert = new CertItem();
            cert.name = certData.getName();
            cert.certData = certData;
            cert.cert = new Certificate(Hex.decode(certData.getContent()));
            if (certData.getPrivateKey() != null && !certData.getPrivateKey().isEmpty()) {
                byte[] privKeyBytes = CryptoService.get().symmetricDecrypt(aesKey,
                        Hex.decode(certData.getPrivateKey()));
                cert.privKey = CryptoService.get().parsePrivateKey(privKeyBytes);
                Arrays.fill(privKeyBytes, (byte) 0);
            }
            // 将证书添加到列表中
            certs.put(cert.name, cert);
            memberIdSnCerts.put(Pair.of(cert.cert.getMemberID(), cert.cert.getSN()), cert);
        }
    }

    /**
     * 保存密钥库至文件。
     *
     * @param file 文件
     * @throws IOException 输入输出异常
     */
    public void saveToFile(File file) throws IOException {
        try {
            // 保存密钥文件
            log.info("保存密钥库文件: {}", file.getAbsolutePath());
            // 生成KeyBoxData
            KeyBoxData keyBoxData = new KeyBoxData();
            keyBoxData.setSalt(Hex.toHexString(salt));
            this.keys.values().forEach(u -> keyBoxData.getKeys().add(u.keyData));
            this.certs.values().forEach(u -> keyBoxData.getCertificates().add(u.certData));
            // 按照名称排序
            Collections.sort(keyBoxData.getKeys(), Ordering.natural().onResultOf(u -> u.getName()));
            Collections.sort(keyBoxData.getCertificates(), Ordering.natural().onResultOf(u -> u.getName()));
            // 将对象序列化XML
            JAXBContext context = JAXBContext.newInstance(KeyBoxData.class);
            Marshaller marshaller = context.createMarshaller();
            marshaller.setProperty(Marshaller.JAXB_FORMATTED_OUTPUT, Boolean.TRUE);
            marshaller.marshal(keyBoxData, file);
        } catch (Exception e) {
            throw new IOException(e);
        }
    }

    /**
     * 修改密码。
     *
     * @param newPassword 新密码
     */
    public void changePassword(String newPassword) {
        // 重新生成盐
        this.salt = CryptoService.get().generateKey();
        // 重新生成加密器
        this.aesKey = CryptoService.get().generateKey(newPassword, salt, false);
        // 重新生成序列化密钥数据
        for (ECKeyItem ecKeyItem : keys.values()) {
            ecKeyItem.keyData = new ECKeyData();
            ecKeyItem.keyData.setName(ecKeyItem.name);
            ecKeyItem.keyData.setPublicKey(ecKeyItem.key.getPublicKeyAsHex());
            if (ecKeyItem.key.hasPrivKey()) {
                byte[] privKeyBytes = ecKeyItem.key.getPrivKeyBytes();
                byte[] encoded = CryptoService.get().symmetricEncrypt(aesKey, privKeyBytes);
                Arrays.fill(privKeyBytes, (byte) 0);
                ecKeyItem.keyData.setPrivateKey(Hex.toHexString(encoded));
            }
        }
        // 重新生成序列化证书数据
        for (CertItem certItem : certs.values()) {
            certItem.certData = new CertificateData();
            certItem.certData.setName(certItem.name);
            certItem.certData.setContent(Hex.toHexString(certItem.cert.getBytes()));
            if (certItem.privKey != null) {
                byte[] privKeyBytes = CryptoUtils.getECPrivateKeyBytes(certItem.privKey);
                byte[] encoded = CryptoService.get().symmetricEncrypt(aesKey, privKeyBytes);
                Arrays.fill(privKeyBytes, (byte) 0);
                certItem.certData.setPrivateKey(Hex.toHexString(encoded));
            }
        }
    }

    /**
     * 添加密钥。
     *
     * @param name 密钥名称
     * @param key  密钥
     */
    public void addKey(String name, ECKey key) {
        ECKeyItem item = new ECKeyItem();
        item.name = name;
        item.key = key;
        item.keyData = new ECKeyData();
        item.keyData.setName(name);
        item.keyData.setPublicKey(key.getPublicKeyAsHex());
        if (key.hasPrivKey()) {
            byte[] privKeyBytes = key.getPrivKeyBytes();
            byte[] encoded = CryptoService.get().symmetricEncrypt(aesKey, privKeyBytes);
            Arrays.fill(privKeyBytes, (byte) 0);
            item.keyData.setPrivateKey(Hex.toHexString(encoded));
        }
        this.keys.put(name, item);
    }

    /**
     * 添加证书。
     *
     * @param name        证书名称
     * @param certificate 证书实例
     * @param privateKey  证书对应的私钥
     */
    public void addCertificate(String name, Certificate certificate, PrivateKey privateKey) {
        CertItem item = new CertItem();
        item.name = name;
        item.privKey = privateKey;
        item.cert = certificate;
        item.certData = new CertificateData();
        item.certData.setName(name);
        item.certData.setContent(Hex.toHexString(certificate.getBytes()));
        if (privateKey != null) {
            byte[] privKeyBytes = CryptoUtils.getECPrivateKeyBytes(privateKey);
            byte[] encoded = CryptoService.get().symmetricEncrypt(aesKey, privKeyBytes);
            Arrays.fill(privKeyBytes, (byte) 0);
            item.certData.setPrivateKey(Hex.toHexString(encoded));
        }
        this.certs.put(name, item);
        this.memberIdSnCerts.put(Pair.of(certificate.getMemberID(), certificate.getSN()), item);
    }

    /**
     * 添加证书。
     *
     * @param name        证书名称
     * @param certificate 证书实例
     */
    public void addCertificate(String name, Certificate certificate) {
        addCertificate(name, certificate, null);
    }

    /**
     * 根据MEMBER ID获取证书信息。
     *
     * @param memberID 用户ID
     * @return 返回证书信息和所对应的私钥
     */
    public List<Pair<Certificate, PrivateKey>> getCertificates(String memberID) {
        return memberIdSnCerts.entrySet().stream().filter(e -> e.getKey().getLeft().equals(memberID))
                .map(e -> Pair.of(e.getValue().cert, e.getValue().privKey)).collect(Collectors.toList());
    }

    /**
     * 根据SN获取证书信息。
     *
     * @param sn 序列号
     * @return 返回证书信息和所对应的私钥
     */
    public List<Pair<Certificate, PrivateKey>> getCertificates(BigInteger sn) {
        return memberIdSnCerts.entrySet().stream().filter(e -> e.getKey().getRight().equals(sn))
                .map(e -> Pair.of(e.getValue().cert, e.getValue().privKey)).collect(Collectors.toList());
    }

    /**
     * 根据MEMBER ID，SN获取证书信息
     *
     * @param memberID 用户ID
     * @param sn       序列号
     * @return 返回证书信息和所对应的私钥
     */
    public Pair<Certificate, PrivateKey> getCertificate(String memberID, BigInteger sn) {
        Pair<String, BigInteger> mapKey = Pair.of(memberID, sn);
        CertItem item = this.memberIdSnCerts.get(mapKey);
        if (item == null)
            return null;
        else
            return Pair.of(item.cert, item.privKey);
    }

    /**
     * 返回所有密钥。
     *
     * @return 所有密钥的map
     */
    public Map<String, ECKey> getKeys() {
        Map<String, ECKey> ret = new HashMap<>();
        keys.entrySet().forEach(e -> ret.put(e.getKey(), e.getValue().key));
        return ret;
    }

    /**
     * 返回所有证书。
     *
     * @return 所有证书的map
     */
    public Map<String, Pair<Certificate, PrivateKey>> getCertificates() {
        Map<String, Pair<Certificate, PrivateKey>> ret = new HashMap<>();
        certs.entrySet().forEach(e -> ret.put(e.getKey(), Pair.of(e.getValue().cert, e.getValue().privKey)));
        return ret;
    }

    /**
     * 获取密钥。
     *
     * @param name 密钥名称
     * @return 密钥
     */
    public ECKey getKey(String name) {
        ECKeyItem key = keys.get(name);
        return key != null ? key.key : null;
    }

    /**
     * 获取证书和对应的私钥。
     *
     * @param name 证书名称
     * @return 证书和对应的私钥
     */
    public Pair<Certificate, PrivateKey> getCertWithPrivKey(String name) {
        CertItem cert = certs.get(name);
        return cert != null ? Pair.of(cert.cert, cert.privKey) : null;
    }

    /**
     * 获取证书。
     *
     * @param name 证书名称
     * @return 证书
     */
    public Certificate getCertificate(String name) {
        CertItem cert = certs.get(name);
        return cert != null ? cert.cert : null;
    }

    /**
     * 获取证书对应的私钥。
     *
     * @param name 证书名称
     * @return 私钥
     */
    public PrivateKey getCertificatePrivateKey(String name) {
        CertItem cert = certs.get(name);
        return cert != null ? cert.privKey : null;
    }

    /**
     * 重命名密钥。
     *
     * @param oldName 旧名称
     * @param newName 新名称
     */
    public void renameKey(String oldName, String newName) {
        ECKeyItem item = keys.remove(oldName);
        if (item == null)
            return;
        item.name = newName;
        item.keyData.setName(newName);
        keys.put(newName, item);
    }

    /**
     * 重命名证书。
     *
     * @param oldName 旧名称
     * @param newName 新名称
     */
    public void renameCertificate(String oldName, String newName) {
        CertItem item = certs.remove(oldName);
        if (item == null)
            return;
        item.name = newName;
        item.certData.setName(newName);
        certs.put(newName, item);
    }

    /**
     * 删除密钥。
     *
     * @param keyName 密钥名称
     */
    public void removeKey(String keyName) {
//        ECKeyItem item = keys.remove(keyName);
        keys.remove(keyName);
    }

    /**
     * 删除证书。
     *
     * @param certName 证书名称
     */
    public void removeCertificate(String certName) {
        CertItem item = certs.remove(certName);
        memberIdSnCerts.remove(Pair.of(item.cert.getMemberID(), item.cert.getSN()));
    }

    /**
     * 检查是否包含密钥。
     *
     * @param name 密钥名称
     * @return 是否包含密钥
     */
    public boolean containsKey(String name) {
        return keys.containsKey(name);
    }

    /**
     * 检查是否包含证书。
     *
     * @param name 证书名称
     * @return 是否包含证书
     */
    public boolean containsCertificate(String name) {
        return certs.containsKey(name);
    }

}
